import akshare as ak
import pandas as pd
import requests
from datetime import datetime
from typing import List, Dict, Optional
from loguru import logger
from config.settings import settings
from util import get_latest_trading_date, get_previous_trading_date

class StockIndexService:
    """股指数据服务类，用于获取股指相关数据"""
    
    def get_index_spot_data(self, symbol_list: List[str],
                           start_datetime: Optional[str] = None,
                           end_datetime: Optional[str] = None,
                           return_raw: bool = False,
                           calculate_change: bool = True) -> Optional[pd.DataFrame | List[Dict]]:
        """
        获取指数分钟级数据（最新数据可作为实时行情）
        
        Args:
            symbol_list: 指数代码列表，如["000300", "000905"]，必须由外部传入
            start_datetime: 开始时间，格式为YYYY-MM-DD HH:MM:SS，如果为None且calculate_change为True则自动获取前一交易日和当日数据
            end_datetime: 结束时间，格式为YYYY-MM-DD HH:MM:SS
            return_raw: 是否返回原始DataFrame，默认False返回DataFrame（按code排序）
            calculate_change: 是否计算涨跌额和涨跌幅，默认True
            
        Returns:
            返回DataFrame（字段名为英文），按code排序；如果return_raw为True，返回未处理的原始数据
            如果calculate_change为True，则返回包含change和change_pct字段的数据
            如果获取失败返回 None
        """
        try:
            # 如果需要计算涨跌幅且没有指定时间范围，自动设置时间范围
            if calculate_change and start_datetime is None:
                # 获取前一交易日和当日交易日
                previous_trading_date = get_previous_trading_date()
                latest_trading_date = get_latest_trading_date()
                
                if previous_trading_date and latest_trading_date:
                    start_datetime = f"{previous_trading_date} 09:30:00"
                    if end_datetime is None:
                        end_datetime = f"{latest_trading_date} 15:00:00"
                    logger.info(f"Auto set time range: {start_datetime} to {end_datetime}")
                else:
                    logger.warning("Failed to get trading dates, using default time range")
                    # 如果获取交易日失败，回退到当前日期
                    today = datetime.now()
                    start_datetime = today.strftime("%Y-%m-%d 09:30:00")
                    if end_datetime is None:
                        end_datetime = today.strftime("%Y-%m-%d 15:00:00")
            
            # 构建请求参数
            params = {"codes": symbol_list}
            if start_datetime:
                params["start_datetime"] = start_datetime
            if end_datetime:
                params["end_datetime"] = end_datetime
            
            # 构建API URL
            url = f"{settings.data_query_host}/index/minute"
            
            logger.info(f"Requesting index minute data from /index/minute with params: {params}")
            
            # 发送HTTP请求
            response = requests.get(
                url,
                params=params,
                timeout=settings.request_timeout_seconds,
            )
            response.raise_for_status()
            
            # 解析响应
            result = response.json()
            
            if result.get("code") != 200:
                logger.error(f"API returned error: {result.get('message', 'Unknown error')}")
                return None
            
            data = result.get("data", [])
            if not data:
                logger.warning("No data returned from API")
                return None
            
            # 转换为DataFrame
            df = pd.DataFrame(data)
            
            # 重命名字段以保持兼容性
            column_mapping = {
                'datetime': 'datetime',
                'code': 'code',
                'index_name': 'name',
                'open': 'open',
                'close': 'close',
                'high': 'high',
                'low': 'low',
                'volume': 'volume',
                'money': 'amount'
            }
            
            # 只重命名存在的列
            existing_columns = {k: v for k, v in column_mapping.items() if k in df.columns}
            df = df.rename(columns=existing_columns)
            
            # 添加price字段作为close的别名，保持向后兼容
            if 'close' in df.columns:
                df['price'] = df['close']
            
            # 如果要求返回原始数据
            if return_raw:
                return df
            
            # 确保datetime列为datetime类型
            if 'datetime' in df.columns:
                df['datetime'] = pd.to_datetime(df['datetime'])
            
            # 确保code列为字符串并填充为6位
            if 'code' in df.columns:
                df['code'] = df['code'].astype(str).str.zfill(6)
            
            # 按code排序
            df = df.sort_values('code').reset_index(drop=True)
            
            # 如果需要计算涨跌幅
            if calculate_change and not return_raw:
                change_df = self.calculate_change_metrics(df)
                if change_df is not None:
                    df = change_df
            
            logger.info(f"Successfully fetched {len(df)} records of index spot data")
            return df
            
        except requests.exceptions.RequestException as e:
            logger.error(f"HTTP request failed: {e}")
            return None
        except Exception as e:
            logger.error(f"Failed to get index spot data: {e}")
            return None
    
    def get_history_data(self, codes: str | List[str], period='1d', start_date='2010-01-01', end_date=None) -> Optional[pd.DataFrame]:
        """
        获取指数历史数据，调用/data_query/index_history_data接口
        
        Args:
            codes (str or List[str]): 指数代码或代码列表
            period (str): 数据周期, '1d', '1w', '1M', '1Y'
            start_date (str): 开始日期
            end_date (str): 结束日期

        Returns:
            pd.DataFrame: 包含历史数据的DataFrame，如果失败返回None
        """
        try:
            # 如果没有提供结束日期，使用当前日期
            if end_date is None:
                end_date = datetime.now().strftime("%Y-%m-%d")

            # 构建请求参数
            params = {
                "start_date": start_date,
                "end_date": end_date,
                "codes": codes
            }
            
            # 构建API URL
            url = f"{settings.data_query_host}/data_query/index_history_data"
            
            logger.info(f"Requesting index history data from /data_query/index_history_data with params: {params}")
            
            # 发送HTTP请求
            response = requests.get(
                url,
                params=params,
                timeout=settings.request_timeout_seconds,
            )
            response.raise_for_status()
            
            # 解析响应
            result = response.json()
            
            if result.get("code") != 200:
                logger.error(f"API returned error: {result.get('message', 'Unknown error')}")
                return None
            
            data = result.get("data", [])
            if not data:
                logger.warning("No data returned from API")
                return None
            
            # 转换为DataFrame
            df = pd.DataFrame(data)
            
            # 确保日期列为datetime类型
            if 'date' in df.columns:
                df['date'] = pd.to_datetime(df['date'])
                if 'code' in df.columns:
                    df = df.sort_values(['code', 'date']).reset_index(drop=True)
                else:
                    df = df.sort_values('date').reset_index(drop=True)

            if period != '1d':
                freq_map = {'1w': 'W-MON', '1M': 'M', '1Y': 'Y'}
                freq = freq_map.get(period)
                if freq:
                    if 'code' in df.columns and not df.empty:
                        # 使用 groupby 和 apply 进行分组重采样
                        df = df.groupby('code', group_keys=False).apply(self._resample_data, freq=freq)
                    elif not df.empty:
                        df = self._resample_data(df, freq)
            
            logger.info(f"Successfully fetched {len(df)} records of index history data")
            return df
            
        except requests.exceptions.RequestException as e:
            logger.error(f"HTTP request failed: {e}")
            return None
        except Exception as e:
            logger.error(f"Failed to get index history data: {e}")
            return None

    def _resample_data(self, df: pd.DataFrame, freq: str) -> pd.DataFrame:
        """
        通用数据重采样方法

        Args:
            df: 包含日数据的DataFrame，必须包含date列和OHLCV数据
            freq: 重采样频率 ('W-MON', 'M', 'Y')
            
        Returns:
            重采样后的DataFrame
        """
        if df.empty or 'date' not in df.columns:
            return df
            
        df_copy = df.copy().set_index('date')
        
        # 定义重采样规则
        agg_rules = {}
        if 'open' in df_copy.columns:
            agg_rules['open'] = 'first'
        if 'high' in df_copy.columns:
            agg_rules['high'] = 'max'
        if 'low' in df_copy.columns:
            agg_rules['low'] = 'min'
        if 'close' in df_copy.columns:
            agg_rules['close'] = 'last'
        if 'volume' in df_copy.columns:
            agg_rules['volume'] = 'sum'
        if 'money' in df_copy.columns:
            agg_rules['money'] = 'sum'
            
        for col in df_copy.columns:
            if col not in agg_rules and df_copy[col].dtype in ['float64', 'int64']:
                agg_rules[col] = 'last'
                
        resampled_df = df_copy.resample(freq).agg(agg_rules)
        resampled_df.reset_index(inplace=True)
        resampled_df.dropna(subset=['open'] if 'open' in resampled_df.columns else resampled_df.columns[:1], inplace=True)
        
        return resampled_df.sort_values('date').reset_index(drop=True)

    def get_index_minute_data(self, codes: str | List[str], 
                             start_datetime: Optional[str] = None,
                             end_datetime: Optional[str] = None) -> Optional[pd.DataFrame]:
        """
        获取指数分钟级数据，调用/data_query/index_minute_data接口
        
        Args:
            codes (str or List[str]): 指数代码或代码列表
            start_datetime (str, optional): 开始时间，格式为YYYY-MM-DD HH:MM:SS
            end_datetime (str, optional): 结束时间，格式为YYYY-MM-DD HH:MM:SS
            
        Returns:
            pd.DataFrame: 包含分钟级数据的DataFrame，如果失败返回None
        """
        try:
            # 构建请求参数
            params = {"codes": codes}
            if start_datetime:
                params["start_datetime"] = start_datetime
            if end_datetime:
                params["end_datetime"] = end_datetime
            
            # 构建API URL
            url = f"{settings.data_query_host}/data_query/index_minute_data"
            
            logger.info(f"Requesting index minute data from /data_query/index_minute_data with params: {params}")
            
            # 发送HTTP请求
            response = requests.get(
                url,
                params=params,
                timeout=settings.request_timeout_seconds,
            )
            response.raise_for_status()
            
            # 解析响应
            result = response.json()
            
            if result.get("code") != 200:
                logger.error(f"API returned error: {result.get('message', 'Unknown error')}")
                return None
            
            data = result.get("data", [])
            if not data:
                logger.warning("No minute data returned from API")
                return None
            
            # 转换为DataFrame
            df = pd.DataFrame(data)
            
            # 确保时间列为datetime类型
            if 'datetime' in df.columns:
                df['datetime'] = pd.to_datetime(df['datetime'])
                # 按code和时间排序
                if 'code' in df.columns:
                    df = df.sort_values(['code', 'datetime']).reset_index(drop=True)
                else:
                    df = df.sort_values('datetime').reset_index(drop=True)
            
            logger.info(f"Successfully fetched {len(df)} records of index minute data")
            return df
            
        except requests.exceptions.RequestException as e:
            logger.error(f"HTTP request failed: {e}")
            return None
        except Exception as e:
            logger.error(f"Failed to get index minute data: {e}")
            return None

    def calculate_change_metrics(self, df: pd.DataFrame) -> Optional[pd.DataFrame]:
        """
        计算涨跌额和涨跌幅（向量化计算）
        分钟级数据的涨跌幅与前一天的收盘价比较
        
        Args:
            df: 包含多天数据的DataFrame
            
        Returns:
            包含最新价格、涨跌额、涨跌幅的DataFrame
        """
        try:
            if df.empty or 'datetime' not in df.columns or 'price' not in df.columns:
                return None
            
            # 确保datetime列为datetime类型
            df['datetime'] = pd.to_datetime(df['datetime'])
            
            # 添加日期列用于分组
            df['date'] = df['datetime'].dt.date
            
            # 按代码和时间排序
            df = df.sort_values(['code', 'datetime']).reset_index(drop=True)
            
            # 获取每个代码每天的最后一条记录作为当天收盘价
            daily_close = df.groupby(['code', 'date']).agg({
                'price': 'last',
                'name': 'first'
            }).reset_index()
            
            # 按代码和日期排序，然后使用shift获取前一个交易日的收盘价
            daily_close = daily_close.sort_values(['code', 'date']).reset_index(drop=True)
            daily_close['prev_close'] = daily_close.groupby('code')['price'].shift(1)
            
            # 合并前一天收盘价到原数据
            df = df.merge(
                daily_close[['code', 'date', 'prev_close']],
                on=['code', 'date'],
                how='left'
            )
            
            # 获取每个代码的最新记录
            latest_df = df.groupby('code').tail(1).copy()
            
            # 计算涨跌额和涨跌幅
            latest_df['change'] = latest_df['price'] - latest_df['prev_close']
            latest_df['change_pct'] = ((latest_df['change'] / latest_df['prev_close']) * 100).fillna(0)
            
            # 对于没有前一天收盘价的记录，设置涨跌为0
            latest_df.loc[latest_df['prev_close'].isna(), 'change'] = 0
            latest_df.loc[latest_df['prev_close'].isna(), 'change_pct'] = 0
            
            # 选择需要的字段
            result_columns = ['code', 'name', 'price', 'change', 'change_pct', 'datetime']
            
            # 如果有close字段也保留
            if 'close' in latest_df.columns:
                result_columns.append('close')
            
            # 添加其他可能需要的字段
            for col in ['open', 'high', 'low', 'volume', 'amount']:
                if col in latest_df.columns:
                    result_columns.append(col)
            
            result_df = latest_df[result_columns].reset_index(drop=True)
            
            return result_df
                
        except Exception as e:
            logger.error(f"Failed to calculate change metrics: {e}")
            return None

